from torch.utils.data import Dataset
from utils.utils import get_source_list
from utils.utils import truc_tensor
import numpy as np
import torch
from typing import Optional


class SingleDataset(Dataset):
    def __init__(self, scp_path: str, length: Optional[int], typ="npy", rank=-1):
        """
        Initialize single dataset using scp file,
        scp_path: file path, where the files should contain each source path in each line
        length: the data length
        type: file type, [npy]
        """
        self.scp_list = get_source_list(scp_path)
        self.typ = typ
        self.length = length

    def __len__(self):
        return len(self.scp_list)

    def __getitem__(self, idx):
        path = self.scp_list[idx]
        if self.typ == "npy":
            res = truc_tensor(torch.from_numpy(np.load(path)), length=self.length)
            return res, path
        else:
            raise Exception("Not implemented")
